package com.imooc.springcloud.rules;

import com.netflix.client.config.IClientConfig;
import com.netflix.loadbalancer.AbstractLoadBalancerRule;
import com.netflix.loadbalancer.IRule;
import com.netflix.loadbalancer.Server;
import lombok.NoArgsConstructor;
import org.springframework.util.CollectionUtils;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;
import org.springframework.web.servlet.support.RequestContext;

import javax.servlet.http.HttpServletRequest;
import java.io.UnsupportedEncodingException;
import java.nio.charset.StandardCharsets;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.Collections;
import java.util.List;
import java.util.SortedMap;
import java.util.TreeMap;

/**
 * @describte:自定义IRule，使用一致性哈希
 * @program:spring-cloud-demo
 * @author:zhengdp
 * @create:2021-10-01 16:28:52
 **/
@NoArgsConstructor
public class MyRule extends AbstractLoadBalancerRule implements IRule {
    @Override
    public void initWithNiwsConfig(IClientConfig clientConfig) {
    }

    /**
     *
     * @param key 默认值为default
     * @return
     */
    @Override
    public Server choose(Object key) {
        HttpServletRequest request = ((ServletRequestAttributes)
                                        RequestContextHolder.getRequestAttributes())
                                        .getRequest();
        // 使用 url作为摘要
        // ServletPath  host 后面的路径
        String uri = request.getServletPath() + "?" + request.getQueryString();
        return route(uri.hashCode(),getLoadBalancer().getAllServers());
    }

    /**
     * 根据
     * @param hashId
     * @param addressList
     * @return
     */
    public Server route(int hashId, List<Server> addressList){
        if(CollectionUtils.isEmpty(addressList)){
            return null;
        }
        // 排序 结点 --> 在环上的位置
        TreeMap<Long,Server> address = new TreeMap<>();

        addressList.stream().forEach(e -> {
            // 虚化 ， 将一个服务结点 虚化成若干个服务结点到环上
            for(int i =0 ; i < 8 ; i++){
                long hash  = hash(e.getId() + i);
                address.put(hash,e);
            }
        });

        long hash = hash(String.valueOf(hashId));
        // 顺时针方式找到服务
        SortedMap<Long,Server> last = address.tailMap(hash);
        // 当 RequestU URL 的Hash值 大于所有的服务结点对应的Hashkey
        // 取address中的第一个结点 （ 因为一致性哈希是一个首尾相接的环）
        if(last.isEmpty()){
            // 没有比 hash 大的服务节点
            return address.firstEntry().getValue();
        }
        return last.get(last.firstKey());
    }

    public long hash(String key){
        MessageDigest md5;
        try{
            md5 = MessageDigest.getInstance("MD5");
        } catch (NoSuchAlgorithmException e) {
            e.printStackTrace();
            throw new RuntimeException(e);
        }
        byte[] keyByte = null;
        try {
            keyByte = key.getBytes("UTF-8");
        } catch (UnsupportedEncodingException e) {
            e.printStackTrace();
            throw new RuntimeException(e);
        }

        md5.update(keyByte);
        byte[] digest = md5.digest(); // 利用md5 生成摘要
        long hashCode = ((long) (digest[2] & 0xFF << 16))
                        | ((long) (digest[1] & 0xFF << 16))
                        | ((long) (digest[0] & 0xFF));
        return hashCode & 0xffffffff;
    }
}
